# -*- coding: utf-8 -*-
import datetime
from pymongo import MongoClient
from zc_core.pipelines.box import BoxPipeline
from zc_core.util.batch_gen import batch_to_year


class PlapBoxPipeline(BoxPipeline):

    def open_spider(self, spider):
        _ = spider
        self.client = MongoClient(self.mongo_uri)
        self.customer = spider.customer
        if not self.customer:
            raise Exception("客户编码未配置： [{}]".format(spider))
        # 默认初始化当前年的库
        year = str(datetime.datetime.now().year)
        self.db_map[year] = self.client['{}_{}_{}'.format(self.bot_name, self.customer, year)]

    def get_db(self, batch_no=None, year=None):
        if not batch_no and not year:
            raise Exception('批次编号与年份至少指定一个')
        if not year:
            year = batch_to_year(batch_no)
        db = self.db_map.get(year)
        if not db:
            db = self.client['{}_{}_{}'.format(self.bot_name, self.customer, year)]
            self.db_map[year] = db

        return db
